from tests.pytests.helpers.codegen_helpers import compile_chunk, assert_code_contains


def test_augassign():
    def g(n: int) -> int:
        s = 0
        i = 0
        while i < n:
            s += i
            i += 1
        return s

    block = compile_chunk(g)
    assert_code_contains(
        block,
        """
        int g(int n) {
            int s = 0;
            int i = 0;
            while (i < n) {
                s += i;
                i += 1;
            }
            return s;
        }
        """,
    )


def test_multi_target_assignment():
    """测试多目标赋值：a = b = value"""

    def multi_assign() -> int:
        a = b = c = 10
        a += 1
        b *= 2
        return a + b + c

    block = compile_chunk(multi_assign)
    assert_code_contains(
        block,
        """
        int multi_assign() {
        int c = 10;
        int b = c;
        int a = b;
        a += 1;
        b *= 2;
        return a + b + c;
        }
        """,
    )


def test_tuple_unpacking_assignment():
    """测试元组解包赋值：a, b = value"""

    def tuple_unpack() -> int:
        x, y = 5, 10
        x += y
        y -= 2
        return x + y

    block = compile_chunk(tuple_unpack)
    assert_code_contains(
        block,
        """
        int tuple_unpack() {
            int x = 5;
            int y = 10;
            x += y;
            y -= 2;
            return x + y;
        }
        """,
    )


def test_complex_augmented_assignment():
    """测试复杂的增强赋值操作"""

    def complex_aug() -> int:
        result = 100
        # 测试各种增强赋值操作符
        result += 10  # 加法
        result -= 5  # 减法
        result *= 2  # 乘法
        result //= 3  # 整除
        result %= 7  # 取模
        result **= 2  # 幂运算
        return result

    block = compile_chunk(complex_aug)
    assert_code_contains(
        block,
        """
        int complex_aug() {
            int result = 100;
            result += 10;
            result -= 5;
            result *= 2;
            result /= 3;
            result %= 7;
            result = mc::pow(result, 2);
            return result;
        }
        """,
    )


def test_nested_assignment_with_augassign():
    """测试嵌套赋值与增强赋值的组合"""

    def nested_assign() -> int:
        # 多层嵌套的变量声明和增强赋值
        outer = 0
        if True:
            inner = 5
            outer += inner
            if True:
                nested = 3
                inner *= nested
                outer += nested
        return outer

    block = compile_chunk(nested_assign)
    assert_code_contains(
        block,
        """
        int nested_assign() {
            int outer = 0;
            if (true) {
                int inner = 5;
                outer += inner;
                if (true) {
                    int nested = 3;
                    inner *= nested;
                    outer += nested;
                }
            }
            return outer;
        }
        """,
    )


def test_assignment_with_function_calls():
    """测试赋值与函数调用的组合"""

    src_code = """
def helper(x: int) -> int:
    return x * 2

def assign_with_calls() -> int:
    a = helper(5)
    b = c = helper(a)
    a += helper(b)
    return a + b + c
"""
    block = compile_chunk(src_code)
    assert_code_contains(
        block,
        """
        int assign_with_calls() {
            int a = helper(5);
            int c = helper(a);
            int b = c;
            a += helper(b);
            return a + b + c;
        }
        """,
    )


def test_global_variable_assignment():
    """测试 global 变量赋值"""
    # 模块级变量

    src_code = """
global_counter = 0

def increment_global() -> int:
    global global_counter
    global_counter += 10
    local_var = global_counter
    local_var *= 2
    return local_var + global_counter
"""
    block = compile_chunk(src_code)
    assert_code_contains(
        block,
        """
        int increment_global() {
            global_counter += 10;
            int local_var = global_counter;
            local_var *= 2;
            return local_var + global_counter;
        }
        """,
    )


def test_mixed_assignment_patterns():
    """测试混合赋值模式"""

    def mixed_patterns() -> int:
        # 单赋值
        a = 5

        # 多目标赋值
        b = c = a * 2

        # 解包赋值
        x, y = b, c

        # 增强赋值
        a += x
        b *= y

        # 复杂表达式赋值
        result = a + b - (x * y)
        result **= 2

        return result

    block = compile_chunk(mixed_patterns)
    assert_code_contains(
        block,
        """
        int mixed_patterns() {
            int a = 5;
            int c = a * 2;
            int b = c;
            int x = b;
            int y = c;
            a += x;
            b *= y;
            int result = a + b - x * y;
            result = mc::pow(result, 2);
            return result;
        }
        """,
    )


def test_chained_augmented_assignment():
    """测试链式增强赋值"""

    def chained_aug() -> int:
        a = b = c = 10

        # 链式增强赋值
        a += 5
        b += a
        c += b

        # 混合操作
        a *= 2
        b //= 3
        c %= 7

        return a + b + c

    block = compile_chunk(chained_aug)
    assert_code_contains(
        block,
        """
        int chained_aug() {
            int c = 10;
            int b = c;
            int a = b;
            a += 5;
            b += a;
            c += b;
            a *= 2;
            b /= 3;
            c %= 7;
            return a + b + c;
        }
        """,
    )


def test_variable_unpacking_assignment():
    """测试变量解包赋值"""

    def variable_unpack() -> int:
        # 元组变量解包
        a = (1, 2)
        c, d = a

        # 列表变量解包
        b = [3, 4]
        e, f = b

        return c + d + e + f

    block = compile_chunk(variable_unpack)
    assert_code_contains(
        block,
        """
        int variable_unpack() {
            std::tuple<int, int> a = std::make_tuple(1, 2);
            int c = mcpy::unpack_get<0>(a);
            int d = mcpy::unpack_get<1>(a);
            mc::array<int> b = {3, 4};
            int e = mcpy::unpack_get<0>(b);
            int f = mcpy::unpack_get<1>(b);
            return c + d + e + f;
        }
        """,
    )


def test_complex_variable_declarations():
    """测试复杂的变量声明场景"""

    def complex_declarations() -> int:
        # 基本变量声明
        x = 10

        # 条件分支中的变量声明
        if x > 5:
            y = 20
            if y > 15:
                z = 30
                x += z
            y += x

        # 循环中的变量声明
        for i in range(3):
            loop_var = i * 2
            x += loop_var

        # 最终结果计算
        result = x * 2
        return result

    block = compile_chunk(complex_declarations)
    assert_code_contains(
        block,
        """
        int complex_declarations() {
            int x = 10;
            if (x > 5) {
                int y = 20;
                if (y > 15) {
                    int z = 30;
                    x += z;
                }
                y += x;
            }
            for (int i = 0; i < 3; ++i) {
                int loop_var = i * 2;
                x += loop_var;
            }
            int result = x * 2;
            return result;
        }
        """,
    )


def test_global_statements_comprehensive():
    """测试全面的 global 语句场景"""

    src_code = """
# 模块级变量
module_var1 = 100
module_var2 = 200
module_var3 = 300

def test_multiple_globals() -> int:
    global module_var1, module_var2
    
    # 修改全局变量
    module_var1 += 50
    module_var2 *= 2
    
    # 局部变量
    local_var = module_var1 + module_var2
    
    return local_var

def test_mixed_global_local() -> int:
    global module_var3
    
    # 全局变量操作
    module_var3 -= 100
    
    # 局部变量声明
    temp = module_var3 * 2
    another_local = temp + 10
    
    # 嵌套作用域
    if temp > 0:
        nested_local = temp / 2
        module_var3 += nested_local
    
    return module_var3 + another_local

def test_global_in_nested_scope() -> int:
    outer_local = 10
    
    if True:
        global module_var1
        inner_local = 20
        module_var1 += inner_local + outer_local
        
        if True:
            deep_local = 30
            module_var1 += deep_local
    
    return module_var1 + outer_local
"""

    block = compile_chunk(src_code)
    assert_code_contains(
        block,
        """
        int test_multiple_globals() {
            module_var1 += 50;
            module_var2 *= 2;
            int local_var = module_var1 + module_var2;
            return local_var;
        }
        """,
    )

    assert_code_contains(
        block,
        """
        int test_mixed_global_local() {
            module_var3 -= 100;
            int temp = module_var3 * 2;
            int another_local = temp + 10;
            if (temp > 0) {
                double nested_local = temp / 2;
                module_var3 += nested_local;
            }
            return module_var3 + another_local;
        }
        """,
    )

    assert_code_contains(
        block,
        """
        int test_global_in_nested_scope() {
            int outer_local = 10;
            if (true) {
                int inner_local = 20;
                module_var1 += inner_local + outer_local;
                if (true) {
                    int deep_local = 30;
                    module_var1 += deep_local;
                }
            }
            return module_var1 + outer_local;
        }
        """,
    )


def test_assignment_patterns_comprehensive():
    """测试全面的赋值模式"""

    def assignment_patterns() -> int:
        # 简单赋值
        a = 1

        # 多目标赋值
        b = c = d = 10

        # 元组解包（字面量）
        e, f = 2, 3

        # 元组解包（变量）
        tuple_var = (4, 5)
        g, h = tuple_var

        # 列表解包（字面量）
        i, j = [6, 7]

        # 列表解包（变量）
        list_var = [8, 9]
        k, l = list_var

        # 混合赋值
        p = q = (10, 11)
        r, s = p

        return a + b + c + d + e + f + g + h + i + j + k + l + r + s

    block = compile_chunk(assignment_patterns)
    assert_code_contains(
        block,
        """
        int assignment_patterns() {
            int a = 1;
            int d = 10;
            int c = d;
            int b = c;
            int e = 2;
            int f = 3;
            std::tuple<int, int> tuple_var = std::make_tuple(4, 5);
            int g = mcpy::unpack_get<0>(tuple_var);
            int h = mcpy::unpack_get<1>(tuple_var);
            int i = 6;
            int j = 7;
            mc::array<int> list_var = {8, 9};
            int k = mcpy::unpack_get<0>(list_var);
            int l = mcpy::unpack_get<1>(list_var);
            std::tuple<int, int> q = std::make_tuple(10, 11);
            std::tuple<int, int> p = q;
            int r = mcpy::unpack_get<0>(p);
            int s = mcpy::unpack_get<1>(p);
            return a + b + c + d + e + f + g + h + i + j + k + l + r + s;
        }
        """,
    )


def test_variable_shadowing():
    """测试变量遮蔽场景"""

    def variable_shadowing() -> int:
        x = 10  # 外层变量

        if True:
            x = 20  # 重新赋值，不是声明
            y = 30  # 新变量声明

            if True:
                x = 40  # 重新赋值
                y = 50  # 重新赋值
                z = 60  # 新变量声明

        # 循环中的变量
        for i in range(3):
            loop_x = i  # 新变量声明
            x += loop_x  # 使用外层变量

        return x

    block = compile_chunk(variable_shadowing)
    assert_code_contains(
        block,
        """
        int variable_shadowing() {
            int x = 10;
            if (true) {
                x = 20;
                int y = 30;
                if (true) {
                    x = 40;
                    y = 50;
                    int z = 60;
                }
            }
            for (int i = 0; i < 3; ++i) {
                int loop_x = i;
                x += loop_x;
            }
            return x;
        }
        """,
    )


def test_function_parameters_vs_locals():
    """测试函数参数与局部变量的区别"""

    def func_with_params(param1: int, param2: int) -> int:
        # 参数不需要声明
        param1 += 10
        param2 *= 2

        # 局部变量需要声明
        local1 = param1 + param2
        local2 = local1 * 3

        # 条件中的局部变量
        if param1 > param2:
            conditional_var = param1 - param2
            local1 += conditional_var

        return local1 + local2

    block = compile_chunk(func_with_params)
    assert_code_contains(
        block,
        """
        int func_with_params(int& param1, int& param2) {
            param1 += 10;
            param2 *= 2;
            int local1 = param1 + param2;
            int local2 = local1 * 3;
            if (param1 > param2) {
                int conditional_var = param1 - param2;
                local1 += conditional_var;
            }
            return local1 + local2;
        }
        """,
    )


def test_augmented_assignment_declarations():
    """测试增强赋值中的变量声明"""

    def augmented_declarations() -> int:
        # 基本变量声明
        counter = 0

        # 增强赋值（不是声明）
        counter += 5
        counter *= 2

        # 条件中的新变量声明
        if counter > 5:
            increment = 10
            counter += increment  # increment 是新声明的

        # 循环中的变量声明和增强赋值
        for i in range(3):
            multiplier = i + 1  # 新声明
            counter += multiplier

        # 复杂表达式中的变量声明
        temp_result = counter * 2
        final_adjustment = temp_result // 3
        counter += final_adjustment

        return counter

    block = compile_chunk(augmented_declarations)
    assert_code_contains(
        block,
        """
        int augmented_declarations() {
            int counter = 0;
            counter += 5;
            counter *= 2;
            if (counter > 5) {
                int increment = 10;
                counter += increment;
            }
            for (int i = 0; i < 3; ++i) {
                int multiplier = i + 1;
                counter += multiplier;
            }
            int temp_result = counter * 2;
            int final_adjustment = (temp_result / 3);
            counter += final_adjustment;
            return counter;
        }
        """,
    )


def test_class_method_variable_declarations():
    """测试类方法中的变量声明"""

    src_code = """
class TestClass:
    def __init__(self):
        self.instance_var = 100
    
    def method_with_locals(self, param: int) -> int:
        # 局部变量声明
        local_var = param * 2
        
        # 条件中的局部变量
        if local_var > 10:
            temp_var = local_var + 5
            local_var = temp_var
        
        # 返回局部变量
        result = local_var + self.instance_var
        return result
"""

    block = compile_chunk(src_code)
    assert_code_contains(
        block,
        """
        int TestClass::method_with_locals(int param) {
            int local_var = param * 2;
            if (local_var > 10) {
                int temp_var = local_var + 5;
                local_var = temp_var;
            }
            int result = local_var + this->instance_var;
            return result;
        }
        """,
    )


def test_literal_containers_assignment():
    """测试字面值容器的赋值"""

    def literal_containers() -> int:
        # 列表字面值
        int_list = [1, 2, 3, 4]
        float_list = [1.0, 2.5, 3.7]

        # 字典字面值
        int_dict = {"a": 1, "b": 2, "c": 3}
        mixed_dict = {"name": "test", "count": 42, "active": True}

        # 元组字面值
        int_tuple = (10, 20, 30)
        mixed_tuple = ("hello", 123, True)

        # 集合字面值
        int_set = {1, 2, 3, 4, 5}

        # 嵌套容器
        nested_list = [[1, 2], [3, 4], [5, 6]]
        nested_dict = {"outer": {"inner": 42}}

        # 使用容器元素
        first_int = int_list[0]
        dict_value = int_dict["a"]
        tuple_element = int_tuple[1]

        return first_int + dict_value + tuple_element

    block = compile_chunk(literal_containers)
    assert_code_contains(
        block,
        """
        int literal_containers() {
            mc::array<int> int_list = {1, 2, 3, 4};
            mc::array<double> float_list = {1.0, 2.5, 3.7};
            mc::dict int_dict = mc::dict{{"a", 1}, {"b", 2}, {"c", 3}};
            mc::dict mixed_dict = mc::dict{{"name", "test"}, {"count", 42}, {"active", true}};
            std::tuple<int, int, int> int_tuple = std::make_tuple(10, 20, 30);
            std::tuple<std::string, int, bool> mixed_tuple = std::make_tuple("hello", 123, true);
            mcpy::variant_set int_set = {1, 2, 3, 4, 5};
            mc::array<mc::array<int>> nested_list = {{1, 2}, {3, 4}, {5, 6}};
            mc::dict nested_dict = mc::dict{{"outer", mc::dict{{"inner", 42}}}};
            int first_int = mcpy::list_getitem(int_list, 0);
            int dict_value = mcpy::cast<int>(int_dict["a"]);
            int tuple_element = mcpy::tuple_getitem<1, int>(int_tuple);
            return first_int + dict_value + tuple_element;
        }
        """,
    )


def test_function_call_parameter_assignment():
    """测试函数调用参数赋值场景"""

    def helper_function(x: int, y: str) -> int:
        return len(y) + x

    def parameter_assignments() -> int:
        # 位置参数
        result1 = helper_function(5, "hello")

        # 参数中使用表达式
        base_value = 20
        result2 = helper_function(base_value * 2, "expression")

        # 参数中使用容器元素
        values = [1, 2, 3]
        names = {"first": "alpha", "second": "beta"}
        result3 = helper_function(values[0], names["first"])

        return result1 + result2 + result3

    src_code = """
def helper_function(x: int, y: str) -> int:
    return len(y) + x

def parameter_assignments() -> int:
    # 位置参数
    result1 = helper_function(5, "hello")
    
    # 参数中使用表达式
    base_value = 20
    result2 = helper_function(base_value * 2, "expression")
    
    # 参数中使用容器元素
    values = [1, 2, 3]
    names = {"first": "alpha", "second": "beta"}
    result3 = helper_function(values[0], names["first"])
    
    return result1 + result2 + result3
    """

    block = compile_chunk(src_code)
    assert_code_contains(
        block,
        """
        int helper_function(int x, const std::string& y) {
            return y.size() + x;
        }
        
        int parameter_assignments() {
            int result1 = helper_function(5, "hello");
            int base_value = 20;
            int result2 = helper_function(base_value * 2, "expression");
            mc::array<int> values = {1, 2, 3};
            mc::dict names = mc::dict{{"first", "alpha"}, {"second", "beta"}};
            int result3 = helper_function(mcpy::list_getitem(values, 0), mcpy::cast<std::string>(names["first"]));
            return result1 + result2 + result3;
        }
        """,
    )


def test_basic_container_operations():
    """测试基础容器操作的赋值"""

    def basic_operations() -> int:
        # 列表操作
        numbers = [1, 2, 3]
        length = len(numbers)

        # 字典操作
        data = {"count": 42}
        dict_size = len(data)

        # 字符串操作
        text = "hello"
        text_length = len(text)

        # 元素访问
        first_num = numbers[0]
        count_val = data["count"]

        return length + dict_size + text_length + first_num + count_val

    block = compile_chunk(basic_operations)
    assert_code_contains(
        block,
        """
        int basic_operations() {
            mc::array<int> numbers = {1, 2, 3};
            int length = numbers.size();
            mc::dict data = mc::dict{{"count", 42}};
            int dict_size = data.size();
            std::string text = "hello";
            int text_length = text.size();
            int first_num = mcpy::list_getitem(numbers, 0);
            int count_val = mcpy::cast<int>(data["count"]);
            return length + dict_size + text_length + first_num + count_val;
        }
        """,
    )


def test_basic_expressions_assignment():
    """测试基础表达式的赋值"""

    def basic_expressions() -> int:
        # 基础赋值
        x = 10
        result = 100

        # 逻辑表达式赋值
        flag1 = True
        flag2 = False
        combined_flag = flag1 and flag2

        # 比较表达式赋值
        value = 15
        is_greater = value > 10

        # 数学运算赋值
        power_result = x * x  # 简化的平方运算
        mod_result = x % 3

        return result + power_result + mod_result

    block = compile_chunk(basic_expressions)
    assert_code_contains(
        block,
        """
        int basic_expressions() {
            int x = 10;
            int result = 100;
            bool flag1 = true;
            bool flag2 = false;
            bool combined_flag = flag1 && flag2;
            int value = 15;
            bool is_greater = value > 10;
            int power_result = x * x;
            int mod_result = x % 3;
            return result + power_result + mod_result;
        }
        """,
    )


def test_conditional_expression_assignment():
    """测试三元运算符（条件表达式）的赋值"""

    def conditional_expressions() -> int:
        # 基础三元运算符
        x = 10
        result = 100 if x > 5 else 50

        # 嵌套三元运算符
        y = 3
        category = "high" if y > 10 else ("medium" if y > 5 else "low")

        # 三元运算符在表达式中
        z = 20
        final = (result * 2) if z > 15 else (result + 10)

        return result + len(category) + final

    block = compile_chunk(conditional_expressions)
    assert_code_contains(
        block,
        """
        int conditional_expressions() {
            int x = 10;
            int result = (x > 5) ? 100 : 50;
            int y = 3;
            std::string category = (y > 10) ? "high" : (y > 5) ? "medium"
            : "low";
            int z = 20;
            int final = (z > 15) ? result * 2 : result + 10;
            return result + category.size() + final;
        }
        """,
    )
